package com.ruoyi.websocket.handler;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.util.UriTemplate;

import java.io.IOException;
import java.net.URI;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

/**
 * @ClassName ChatMsgWebSocketHandle
 * @Description
 *    1. 每个用户通过 openId 建立连接
 *  * 2. 默认订阅广播，可以主动取消订阅/恢复订阅
 *  * 3. 可以给指定用户的所有连接推送消息
 *  * 4. 可以广播消息给所有已订阅的连接
 *  *
 *  * 消息类型：
 *  * - HEARTBEAT   心跳检测
 *  * - SUBSCRIBE   订阅广播
 *  * - UNSUBSCRIBE 取消订阅广播
 *  * - SEND        给指定用户推送
 *  * - BROADCAST   广播给所有已订阅用户
 * @Author li yang
 * @Date 2025/9/8 16:37
 * @Version 1.0
 */
@Component
public class ChatMsgWebSocketHandle  extends TextWebSocketHandler {

    private static final Logger log = LoggerFactory.getLogger(ChatMsgWebSocketHandle.class);

    /* ===== 支持的消息类型 ===== */
    private static final String T_HEARTBEAT = "HEARTBEAT";
    private static final String T_SUBSCRIBE = "SUBSCRIBE";
    private static final String T_UNSUBSCRIBE = "UNSUBSCRIBE";
    private static final String T_SEND = "SEND";
    private static final String T_BROADCAST = "BROADCAST";

    private static final Set<String> SUPPORTED = new HashSet<>(Arrays.asList(
            T_HEARTBEAT, T_SUBSCRIBE, T_UNSUBSCRIBE, T_SEND, T_BROADCAST
    ));

    // 用户 ID -> 该用户的所有会话（可能一人多端同时在线）
    private final Map<String, List<WebSocketSession>> sessionMap = new ConcurrentHashMap<>();

    // 会话订阅状态：是否接收广播（默认 true）
    private final Map<WebSocketSession, Boolean> sessionSubscribed = new ConcurrentHashMap<>();

    private static final ObjectMapper OM = new ObjectMapper();

    // 用于从 URL 提取 openId，例如 ws://host/webSocket/chat/{openId}
    private static final UriTemplate URI_TEMPLATE = new UriTemplate("/webSocket/chat/{openId}");

    /**
     * 从连接 URI 中提取 openId
     */
    private String getOpenId(WebSocketSession session) {
        URI uri = session.getUri();
        if (uri == null) {
            return null;
        }
        Map<String, String> vars = URI_TEMPLATE.match(uri.getPath());
        return vars.get("openId");
    }

    /**
     * 安全获取 JSON 字段
     */
    private static String str(JsonNode n, String f, String def) {
        JsonNode x = n.path(f);
        return x.isMissingNode() || x.isNull() ? def : x.asText(def);
    }

    /* ===== 生命周期回调 ===== */

    /**
     * 连接建立时：记录会话、默认订阅广播
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String openId = getOpenId(session);
        if (openId == null || openId.trim().isEmpty()) {
            log.warn("连接建立但未提供有效 openId，会话ID: {}", session.getId());
            session.close(CloseStatus.NOT_ACCEPTABLE.withReason("未提供有效用户标识"));
            return;
        }
        // 将会话加入该用户列表
        sessionMap.computeIfAbsent(openId, k -> new CopyOnWriteArrayList<>()).add(session);
        // 默认订阅广播
        sessionSubscribed.put(session, true);
        log.info("用户 {} 连接建立，会话ID: {}", openId, session.getId());
        session.sendMessage(new TextMessage(sys("CONNECTED", "连接已建立")));
    }

    /**
     * 收到客户端消息时的处理逻辑
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String openId = getOpenId(session);
        if (openId == null) {
            session.close(CloseStatus.NOT_ACCEPTABLE.withReason("无效会话"));
            return;
        }

        // 解析 JSON
        JsonNode root;
        try {
            root = OM.readTree(message.getPayload());
        } catch (Exception e) {
            session.sendMessage(new TextMessage(err("E_JSON", "消息格式错误")));
            return;
        }

        String type = str(root, "type", "");
        if (!SUPPORTED.contains(type)) {
            session.sendMessage(new TextMessage(err("E_TYPE", "未知消息类型")));
            return;
        }

        try {
            switch (type) {
                case T_HEARTBEAT: {
                    // 心跳 -> pong
                    session.sendMessage(new TextMessage(sys("HEARTBEAT", "pong")));
                    break;
                }
                case T_SUBSCRIBE: {
                    // 恢复广播订阅
                    sessionSubscribed.put(session, true);
                    session.sendMessage(new TextMessage(sys("SUBSCRIBE_OK", "已订阅广播")));
                    break;
                }
                case T_UNSUBSCRIBE: {
                    // 取消广播订阅
                    sessionSubscribed.put(session, false);
                    session.sendMessage(new TextMessage(sys("UNSUBSCRIBE_OK", "已取消广播订阅")));
                    break;
                }
                case T_SEND: {
                    // 给 target 用户推送消息
                    String target = str(root, "target", "");
                    String content = str(root, "content", "");
                    if (target.isEmpty()) {
                        session.sendMessage(new TextMessage(err("E_ARG", "缺少 target")));
                        break;
                    }
                    sendToUser(target, msg(openId, content, T_SEND));
                    break;
                }
                case T_BROADCAST: {
                    // 广播给所有已订阅的用户
                    String content = str(root, "content", "");
                    broadcast(msg(openId, content, T_BROADCAST));
                    break;
                }
                default:
                    session.sendMessage(new TextMessage(err("E_TYPE", "未知消息类型")));
            }
        } catch (Exception e) {
            log.error("处理消息异常: {}", e.getMessage(), e);
            session.sendMessage(new TextMessage(err("E_RUNTIME", "服务器内部错误")));
        }
    }

    /**
     * 连接关闭时：移除会话、清理订阅状态
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String openId = getOpenId(session);
        if (openId != null) {
            List<WebSocketSession> list = sessionMap.get(openId);
            if (list != null) {
                list.remove(session);
                if (list.isEmpty()) sessionMap.remove(openId);
            }
        }
        sessionSubscribed.remove(session);
        log.info("连接关闭 openId={}, sessionId={}, reason={}", openId, session.getId(), status.getReason());
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /* ===== 推送与广播 ===== */

    /**
     * 给指定用户的所有连接推送
     */
    private void sendToUser(String openId, String payload) {
        List<WebSocketSession> list = sessionMap.get(openId);
        if (list == null || list.isEmpty()) {
            log.debug("用户 {} 无活跃连接，未送达", openId);
            return;
        }
        int ok = 0;
        for (WebSocketSession s : list) {
            if (!s.isOpen()) continue;
            try {
                s.sendMessage(new TextMessage(payload));
                ok++;
            } catch (IOException e) {
                log.error("发送给用户 {} 失败: {}", openId, e.getMessage());
            }
        }
        log.debug("向用户 {} 推送完成：成功 {}/{}", openId, ok, list.size());
    }

    /**
     * 广播消息：仅推送给已订阅广播的连接
     */
    private void broadcast(String payload) {
        int total = 0, ok = 0;
        for (List<WebSocketSession> list : sessionMap.values()) {
            for (WebSocketSession s : list) {
                total++;
                if (!Boolean.TRUE.equals(sessionSubscribed.getOrDefault(s, true))) continue;
                if (!s.isOpen()) continue;
                try {
                    s.sendMessage(new TextMessage(payload));
                    ok++;
                } catch (IOException e) {
                    log.error("广播失败: {}", e.getMessage());
                }
            }
        }
        log.debug("广播完成：成功 {}/{}", ok, total);
    }

    /* ===== Payload 构造工具 ===== */

    /**
     * 普通消息
     */
    private String msg(String from, String content, String type) {
        try {
            Map<String, Object> p = new HashMap<>();
            p.put("type", type);
            p.put("from", from);
            p.put("content", content);
            p.put("timestamp", System.currentTimeMillis());
            return OM.writeValueAsString(p);
        } catch (JsonProcessingException e) {
            return err("E_JSON", "消息格式错误");
        }
    }

    /**
     * 系统消息
     */
    private String sys(String event, String content) {
        try {
            Map<String, Object> p = new HashMap<>();
            p.put("type", "SYSTEM");
            p.put("event", event);
            p.put("content", content);
            p.put("timestamp", System.currentTimeMillis());
            return OM.writeValueAsString(p);
        } catch (JsonProcessingException e) {
            return "{\"error\":\"系统消息格式错误\"}";
        }
    }

    /**
     * 错误消息
     */
    private String err(String code, String content) {
        try {
            Map<String, Object> p = new HashMap<>();
            p.put("type", "ERROR");
            p.put("code", code);
            p.put("content", content);
            p.put("timestamp", System.currentTimeMillis());
            return OM.writeValueAsString(p);
        } catch (JsonProcessingException e) {
            return "{\"error\":\"错误消息格式错误\"}";
        }
    }
}